#!/usr/local/bin/python

# Program Flow:
# ------------------------------------
# 1) Blast TCDB against Genome
# 2) Parse XML output
# 3) Collect EVERY hit, write to FASTA
# 4) Run HMMTOP on Hit File
# 5) Create tuple-list of: [(ID,e,TMS,TCID,acc)]
# 6) Pickle List 5 & offer step-wise access

from Bio.Blast.Applications import NcbiblastpCommandline as blastp
from Bio.Blast import NCBIXML
from Bio import SeqIO

import pickle
import os,sys
import tempfile
import hmmtop
import ProjectBio
import tcdb

os.environ['BIOV_DEBUG'] = 'True'

class expand:
	
	def __init__(self):
		tcdb.use_local()
		self.tcdb = os.environ['HOME']+'/db/tcdb'
		self.target = str
		self.mytargets = {}
		self.expect = (0.001,0.01)
		self.tms = (12,12)
		self.debug = True
		self.outfile = 'selected'
		self.transporters = []
		
	def primary_blast(self):
		if os.path.exists('primary.xml') is True:
			return
		print 'Running BLAST...'
		os.system('makeblastdb -logfile /dev/null -in '+self.target)
		blast = blastp()
		blast.db = self.target
		blast.query = self.tcdb
		blast.outfmt = 5
		blast.out = 'primary.xml'
		blast.evalue = 1
		blast.comp_based_stats = "0"
		blast()
		# Blast is complete, now parse the results
		
	def collect_results(self):
		# [Parse elimination.xml and collect IDs that were hits]
		if os.path.exists('myres.db') is True:
			self.transporters = pickle.load(open('./myres.db','r'))
			return
		print 'Collecting results...'
		results = NCBIXML.parse(open('primary.xml','r'))
		for hit in results:
			tcid,family,acc = ProjectBio.ParseTC(hit.query)
			for desc in hit.descriptions:
				target_id=ProjectBio.ParseDefline(desc.title,True).id
				self.transporters.append([target_id,desc.e,tcid,family,acc])
		pickle.dump(self.transporters,open('./myres.db','wb'))
		return
		# Done, saved transporter IDs from the target DB
		
	def save_transporters(self):
		# [Build FASTA DB containing all potential transporters & HMMTOP it]
		myfasta = './all_transporters.faa'
		if os.path.exists(myfasta) is False:
			targetdb = SeqIO.parse(self.target,'fasta')
			targetdb = SeqIO.to_dict(targetdb)
			handle = open(myfasta,'wb')
			myids = list(set([i[0] for i in self.transporters]))
			for item in myids:
				SeqIO.write(targetdb[item],handle,'fasta')
			handle.close()
		if os.path.exists('myres2.db') is True:
			self.transporters = pickle.load(open('myres2.db','r'))
			return
		print 'Calculating TMSs...'
		names = tcdb.Names()
		ht = hmmtop.tools()
		ht.add_library('t',myfasta)
		ht.scan_libraries()
		tmss = ht.results['t']
		newtrans = []
		for item in self.transporters:
			abr = names.get_family_abr(item[2])
			del(item[2])
			item[2:2]=[abr]
			tmc = len(tmss[item[0]]) if item[0] in tmss else 0
			item.append(tmc)
			newtrans.append(item)
		self.transporters = newtrans
		pickle.dump(self.transporters,open('myres2.db','wb'))		
				
	def write_selected(self):
		# [Write TAB file with specific settings]
		selection = [i for i in self.transporters if i[1]>=self.expect[0] and i[1]<=self.expect[1] and i[-1]>=self.tms[0] and i[-1]<=self.tms[1]]
		selection.sort(key=lambda x:ProjectBio.sort_tcid(x[-3]),reverse=False)
		handle = open(self.outfile+'.xls','wb')
		handle.write('GeneID\te-val\tFamily\tTCID\tAcc\tTMSs\n')
		print 'Writing Output...'
		for line in selection:
			myline = "\t".join([str(i) for i in line])
			handle.write(myline+'\n')
		return
		print 'Done, your file has been written to this directory.'

		
if __name__=='__main__':		
	e = expand()
	try:
		e.target = sys.argv[1]
	except:
		print 'Usage: expandtcdb.py <PROTEOME_FILE>\n'
		print 'You will be presented with options later to extract specific transporters'
		quit()
	print 'Calculating preliminary stats. This only needs to be done once, as long as we use the same directory..'
	e.primary_blast()
	e.collect_results()
	e.save_transporters()
	emin = raw_input('Best e-value to tolerate? (Smaller number): ')
	emax = raw_input('Worst e-value to tolerate? (Bigger number): ')
	tmsmin = raw_input('Minimum # of TMSs: ')
	tmsmax = raw_input('Maximum # of TMSs: ')
	e.outfile = raw_input('Filename to write results (no extention): ')
	e.expect = (float(emin),float(emax))
	e.tms = (int(tmsmin),int(tmsmax))
	e.write_selected()